package ants_utils

import (
	"github.com/panjf2000/ants/v2"
	"sync/atomic"
)

// Future 接口
type Future interface {
	Get() (interface{}, error)
	IsDone() bool
}

// ChannelFuture 基于 channel 的实现
type ChannelFuture struct {
	result   TaskResult
	resultCh chan TaskResult
	done     atomic.Value
}

func NewChannelFuture(fn TaskFn) *ChannelFuture {
	future := &ChannelFuture{
		resultCh: make(chan TaskResult, 1),
	}
	future.done.Store(false)

	go func() {
		data, err := fn()
		future.resultCh <- TaskResult{
			Data: data,
			Err:  err,
		}
	}()

	return future
}

func NewChannelFutureWithAnts(fn TaskFn, pool *ants.Pool) *ChannelFuture {
	future := &ChannelFuture{
		resultCh: make(chan TaskResult, 1),
	}
	future.done.Store(false)

	pool.Submit(func() {
		data, err := fn()
		future.resultCh <- TaskResult{
			Data: data,
			Err:  err,
		}
	})

	return future
}

func (f *ChannelFuture) Get() (interface{}, error) {
	if d, _ := f.done.Load().(bool); !d {
		f.result = <-f.resultCh
		f.done.Store(true)
	}

	return f.result.Data, f.result.Err
}

func (f *ChannelFuture) IsDone() bool {
	if d, _ := f.done.Load().(bool); !d {
		select {
		case <-f.resultCh:
			f.done.Store(true)
			return true
		default:
			return false
		}
	}

	return true
}
